import datetime

from sqlalchemy import Column, Integer, String, DateTime, Float, Index, PrimaryKeyConstraint
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.sql import func

# 创建基础模型类
Base = declarative_base()


class HydrometricDataTimeseries(Base):
    """时序数据表 - 使用TimescaleDB优化"""
    __tablename__ = 'hydrometric_data_timeseries'

    id = Column(Integer, autoincrement=True)
    station_id = Column(String(50), nullable=False, comment='测站ID')
    timestamp = Column(DateTime, nullable=False, default=datetime.datetime.utcnow, comment='观测时间')
    flow_rate = Column(Float, nullable=False, comment='观测流量 (m³/s)')
    created_at = Column(DateTime, default=func.now(), comment='记录创建时间')

    # TimescaleDB联合主键：station_id作为空间索引，timestamp作为时间分区
    __table_args__ = (
        PrimaryKeyConstraint('station_id', 'timestamp'),
    )

    def __repr__(self):
        return f"<HydrometricDataTimeseries(station_id='{self.station_id}', timestamp='{self.timestamp}', flow_rate={self.flow_rate})>"


class HydrometricDataRegular(Base):
    """普通数据表 - 带索引的常规PostgreSQL表"""
    __tablename__ = 'hydrometric_data_regular'

    id = Column(Integer, primary_key=True, autoincrement=True)
    station_id = Column(String(50), nullable=False, comment='测站ID')
    timestamp = Column(DateTime, nullable=False, default=datetime.datetime.utcnow, comment='观测时间')
    flow_rate = Column(Float, nullable=False, comment='观测流量 (m³/s)')
    created_at = Column(DateTime, default=func.now(), comment='记录创建时间')

    # 创建复合索引以优化查询性能
    __table_args__ = (
        Index('idx_station_timestamp', 'station_id', 'timestamp'),
        Index('idx_timestamp', 'timestamp'),
        Index('idx_station_id', 'station_id'),
    )

    def __repr__(self):
        return f"<HydrometricDataRegular(station_id='{self.station_id}', timestamp='{self.timestamp}', flow_rate={self.flow_rate})>"


# 表创建和TimescaleDB配置的SQL语句 - 性能优化版本
TIMESCALE_SETUP_SQL = """
-- 启用TimescaleDB扩展
CREATE EXTENSION IF NOT EXISTS timescaledb;

-- 将时序表转换为TimescaleDB hypertable - 性能优化配置
-- 使用7天chunk间隔，减少chunk数量，提升查询性能
SELECT create_hypertable('hydrometric_data_timeseries', 'timestamp', 
                        partitioning_column => 'station_id',
                        number_partitions => 4,
                        chunk_time_interval => INTERVAL '7 days',
                        if_not_exists => TRUE);

-- 启用压缩功能以节省存储空间
ALTER TABLE hydrometric_data_timeseries SET (timescaledb.compress = true);
SELECT add_compression_policy('hydrometric_data_timeseries', INTERVAL '30 days', if_not_exists => TRUE);

-- 可即时压缩老 chunk（测试环境先跑）
-- SELECT compress_chunk('_hyper_12_3353_chunk'::regclass);
"""
